library(librarian)
shelf(Seurat,
      tidyseurat,
      harmony,
      ggpubr,
      SingleR,
      tidySummarizedExperiment,
      magrittr,
      tidyverse)

#' Mark doublets in Seurat by Scrublet
#'
#' @param sobj A Seurat object.
#'
#' Need sctkPythonInstallConda() first, and configured python env in Rstudio, or singlecellTK::sctkSelectConda()
#' A tibble of doublet estimation will be printed out in the function
#'
#' @return A Seurat object with added metadata: 'scrublet_call' and 'scrublet_score'
#'
#' @examples 
#' 
#' marked_sobj <- mark_doublets(sobj)
#' sobj <- subset(marked_sobj, scrublet_call == 'Singlet')
mark_doublets <- function(sobj, sample = 'orig.ident') {
  library(singleCellTK)
  selectSCTKConda(envname = 'r-reticulate')
  sobj <- sobj |>
    as.SingleCellExperiment() |>
    runScrublet(sample = sample) |>
    colData() |>
    as_tibble(rownames = "rowname") |>
    select(c(rowname, contains("scrublet"))) |>
    column_to_rownames("rowname") |>
    AddMetaData(sobj, metadata = _)

  dplyr::count(sobj, scrublet_call) |> print()

  sobj
}


#' Mark doublets in Seurat by scDblFinder
#'
#' @param sobj 
#' @param clusters 
#' @param samples 
#'
#' @return
#' @export
#'
#' @examples
mark_doublets2 <- function(sobj, clusters = 'seurat_clusters',
                           samples = 'orig.ident', bplog = T){
  library(scDblFinder)
  library(BiocParallel)
  message('Running scDblFinder...')
  
  n.samples <- sobj |> pull(samples) |> unique() |> length()
  message('Using ', n.samples, ' workers...')
  
  res <- sobj |>
    as.SingleCellExperiment() |>
    scDblFinder(clusters = clusters,
                samples = samples,
                returnType = 'table',
                BPPARAM = MulticoreParam(n.samples, log = bplog)) |>
    as_tibble(rownames = '.cell') |>
    filter(src == 'real') |>
    mutate(.cell = .cell, scdbl.class = class, .keep = 'none')
  
  res |>
    summarise(dbl.rate = sum(scdbl.class == 'doublet') / n()) |>
    print()
  
  res
}


#' Pre-process seurat in one function
#'
#' @param sobj A Seurat object.
#' @param batch Batch variable in the metadat of `sobj`. Default is "orig.ident".
#' @param skip_norm Skip `NormalizeData()` or not. Useful in TPM data. FALSE by default.
#' 
#' A clustered UMAP will be produced in the process.
#' 
#' @return A normalized, scaled, batch-harmonized, clustered Seurat object with PCA and UMAP DimReduc.
#'
#' @examples
#' sobj <- quick_process_seurat(sobj)
quick_process_seurat <- function(sobj,
                                 batch = "orig.ident", leiden = TRUE,
                                 tsne = FALSE, regress.cc = FALSE,
                                 skip_norm = FALSE, pcs = 20, res = .8) {
  start_time <- Sys.time() |> round()
  if (!skip_norm) {
    sobj <- NormalizeData(sobj)
  }
  
  sobj <- sobj |>
    FindVariableFeatures()
  
  sobj <- if (regress.cc) {
    CellCycleScoring(sobj, s.features = cc.genes.updated.2019$s.genes,
                     g2m.features = cc.genes.updated.2019$g2m.genes) |>
      mutate(cc.diff = S.Score - G2M.Score) |>
      ScaleData(vars.to.regress = 'cc.diff')
  } else {
    ScaleData(sobj)
  }
  
  sobj <- sobj |>
    RunPCA(npcs = max(pcs,50))
  
  reduc <- 'pca'
  
  if (length(batch) > 1 || length(unique(pull(sobj,batch))) > 1) {
    sobj <- sobj |>
      RunHarmony(batch, lambda = rep(1,length(batch)))
    
    reduc <- 'harmony'
  }
  
  sobj <- sobj |>
    RunUMAP(reduction = reduc, dims = 1:pcs) |>
    FindNeighbors(reduction = reduc, dims = 1:pcs)
    
  sobj <- if (leiden) {
    FindClusters(sobj, resolution = res, algorithm = 4, random.seed = 1)
  } else {
    FindClusters(sobj, resolution = res)
  }

  if (tsne) {
    sobj <- sobj |> RunTSNE(reduction = reduc, dims = 1:pcs)
  }
  
  tryCatch(DimPlot(sobj, label = TRUE, label.box = TRUE, repel = TRUE,
                         cols = DiscretePalette(36)) |> print(),
           error = function(e){
             warning('DimPlot ran into error. Fallback to geom_bin2d().')
             g1 <- sobj |> 
               ggplot(aes(x = umap_1, y = umap_2, fill = seurat_clusters)) + 
               geom_bin2d(bins = 512) + 
               theme_classic() + 
               scale_fill_manual(values = DiscretePalette(36))
             
             print(g1)
           })
  
  end_time <- Sys.time() |> round()
  
  message("Start time: ", start_time)
  message("End time: ", end_time)
  
  message("Execution time: ", as.duration(end_time - start_time))

  sobj
}

#' Adjust sparse matrix for drop-like seq specific contamination
#'
#' @param toc A sparse matrix of counts (filtered count matrix) from drop-based seq, e.g., 10X.
#' @param tod A sparse matrix of drops (raw count matrix). Will be auto-estimated from `toc` if not provided.
#' @param cluster A named vector, where entries are the cluster IDs and names are cellIDs. If no names are provided, the order is assumed to match the order in `toc`.
#'
#' @return A adjusted sparse matrix
#' @export
#'
#' @examples
#' library(Seurat)
#' sobj <- read_rds("some_seurat.rds")
#' adj_mat <- adjust_soup_solo(GetAssayData(sobj, 'count'), sobj$seurat_clusters)
adjust_soup_solo <- function(toc, tod = NULL, cluster) {
  library(SoupX)
  if (is.null(tod)) {
    soupProf <- data.frame(
    row.names = rownames(toc),
    est = rowSums(toc) / sum(toc),
    counts = rowSums(toc)
  )

  set.seed(42)
  scNoDrops <- soupProf |>
    SoupChannel(toc, calcSoupProfile = FALSE) |>
    setSoupProfile(soupProfile = soupProf) |>
    setClusters(cluster) |>
    autoEstCont() |>
    adjustCounts(roundToInt = TRUE)
  } else {
    set.seed(42)
    scNoDrops <- tod |>
      SoupChannel(toc) |>
      setClusters(cluster) |>
      autoEstCont() |>
      adjustCounts(roundToInt = TRUE)
  }
}

#' Adjust seurat for soup contamination, always assume no raw matrix is provided
#'
#' @param sobjss input seurat object.
#'
#' @return a de-souped seurat object, without metadata
#' @export
#'
#' @examples 
soupx_seurat <- function(sobjss){
  GetAssayData(sobjss, 'count') %>%
    adjust_soup_solo(cluster = sobjss$seurat_clusters) %>%
    CreateSeuratObject(min.cells = 3, min.features = 200, names.field = 2)
}

#' Annotate Seurat by SingleR
#'
#' @param sobj A clustered Seurat object to be annotated.
#' @param ref A ref SingleCellExperiment object, typically from `celldex` package.
#' @param new_label A character. The new annotation of cell type in output.
#' @param fine_label A boolean. If TRUE, will use `ref$label.fine` to run `SingleR()`, otherwise will use `ref$label.main`. FALSE by default.
#'
#' @return A annotated Seurat object.
#'
#' @examples
#' source('path/to/mod_seurat.R')
#' sobj <- quick_process_seurat(sobj) |>
#' mark_cell_type_singler()
mark_cell_type_singler <- function(sobj,
                                   ref,
                                   fine_label = FALSE,
                                   sc_ref = FALSE,
                                   new_label = "singler_label") {
  singler_res <- sobj |>
    GetAssayData() |>
    SingleR(
      ref = ref,
      clusters = sobj$seurat_clusters,
      labels = ifelse(rep_along(along = ref$label.main, fine_label), ref$label.fine, ref$label.main),
      de.method = ifelse(sc_ref, 'wilcox', 'classic'),
      BPPARAM = BiocParallel::MulticoreParam(progressbar = TRUE, tasks = 0)
    )
  
  plotScoreHeatmap(singler_res,
                   show.labels = TRUE,
                   show.pruned = TRUE,
                   show_colnames = TRUE
  ) |> show()
  
  if (new_label %in% colnames(sobj@meta.data)) {
    sobj <- sobj |> select(-all_of(new_label))
    message("Found existing label, overwriting...")
  }
  
  message('Using ', new_label, ' as colname of the assigned cell types...')
  
  singler_label <- tibble(
    singler_label = singler_res$labels,
    seurat_clusters = as.factor(singler_res@rownames)
  ) |>
    set_names(c(new_label, "seurat_clusters")) |>
    left_join(sobj, y = _)
}

#' Add suffix to colnames of the object
#'
#' useful for single cell count matrix to add sample info to cell barcodes. Will connect with '_'. like 'col1' to 'col1_name'.
#' @param x An object with colnames
#' @param suffix A character, the suffix to add
#'
#' @return `x` with its colnames suffixed with `suffix`.
#'
#' @examples
#' head(iris)
#' iris_suffix <- add_name_field(iris, 'my_suffix')
#' head(iris_suffix)
add_name_suffix <- function(x, suffix, sep = '_') {
  colnames(x) <- colnames(x) |>
    map_chr(\(x)str_c(suffix, sep, x))
  x
}

#' Annotate clusters with their unique marker gene
#'
#' @param to_latent An seurat object.
#' @param big_class Character. 
#' Names of new clusters will be '{marker gene}_`big_class`_{old idents}'.
#'
#' @return An modified seurat object with a new column of annotated clusters in metadata.
#' @export
#'
#' @examples
annotate_latents <- function(to_latent, big_class = 'Cell',
                             logfc.thres = .5){
  markers <- to_latent |>
    FindAllMarkers(only.pos = TRUE, logfc.threshold = logfc.thres) |>
    filter(p_val_adj < .05, str_starts(gene, 'LINC|ENSG', negate = TRUE)) |>
    group_by(gene) |>
    slice_max(avg_log2FC) |>
    group_by(cluster) |>
    slice_max(avg_log2FC)
  
  markers |>
    ungroup() |>
    mutate(latent_cluster = str_c(gene, big_class, cluster, sep = '_'),
           seurat_clusters = cluster) |>
    select(latent_cluster, seurat_clusters) |>
    left_join(to_latent, y = _)}

#' Title
#'
#' @param obc 
#'
#' @return
#' @export
#'
#' @examples
infer_cellschat <- function(obc){
  obc |>
    subsetData() |>
    identifyOverExpressedGenes() |>
    identifyOverExpressedInteractions() |>
    computeCommunProb(population.size = TRUE) |>
    aggregateNet() |>
    computeCommunProbPathway() |>
    netAnalysis_computeCentrality()
}

logtpm.mean <- function(x)log1p(mean(expm1(x)))

bill.violin <- function(srtobj, features, group.by, facet.ncol = NULL,
                        pubsize = FALSE, no_x_text = TRUE) {
  gmeta <- srtobj |>
    as_tibble() |>
    select(.cell, {{group.by}})
  
  g1 <- srtobj |>
    get_abundance_sc_long(features) |>
    left_join(gmeta) |>
    ggplot(aes({{group.by}}, .abundance_RNA)) +
    geom_violin(aes(fill = {{group.by}}), scale = 'width') +
    stat_summary(fun = logtpm.mean, geom = 'crossbar',
                 width = .3, color = 'black') +
    facet_wrap(~.feature, scales = 'free_y', ncol = facet.ncol) +
    theme_pubr() +
    labs(y = 'Normalized expression')
  
  if(pubsize) {
    g1 <- g1 + theme_jpub(theme_classic)
  }
  
  if(no_x_text) {
    g1 <- g1 + theme(axis.text.x = element_blank())
  }
  
  g1
}

#' Read GEO supplementary files into single Seurat
#'
#' @param paths Vector of full paths to all files to read.
#' @param name_regex Regex patterns for str_extract() to extract sample names from `paths`. This will be used as orig.ident of the output.
#'
#' @return A Seurat object.
#' @export
#'
#' @examples
#' mex.path <- list.files('mission/FPP/celik24_Efaecalis', 'gz', full.names = T)
#' sobj <- mex.path |> read_geo_supp('(uninf|inf).+_(1|2)')
#' 
read_geo_supp <- function(paths, name_regex, workers = 1){
  mex.name <- paths |>
    str_extract(name_regex) |>
    unique() |>
    str_replace_all('_','-')
  
  message(str_glue('Demultipexing {length(mex.name)} samples...'))
  message('Glimpse of head samples:')
  
  print(head(mex.name))
  
  options(warn = 2)
  if (workers > 1) {
    require(furrr)
    plan('multisession', workers = 8)
    pmap <- future_pmap
    on.exit(plan('sequential'), add = TRUE)
  }
  mex.trio <-
    list(mtx = 'mtx', cells = 'barcodes', features = 'features') |>
    map(\(x)str_subset(paths, x)) |>
    pmap(safely(ReadMtx), .progress = T) |>
    list_transpose()
  
  options(warn = 1)
  if (!every(mex.trio[['error']], is_null)) {
    names(mex.trio[['error']]) <- mex.name
    mex.err <- mex.trio[['error']] |> discard(is_null)
    warning(str_glue('{length(mex.err)} Error(s) found in list:'))
    print(mex.err)
  }
  
  message('\nAppending sample names...')
  mex.trio <- mex.trio[['result']] |>
    map2(mex.name, safely(add_name_suffix)) |>
    list_transpose() |>
    pluck('result') |>
    discard(is_null)
  
  message('Merging matrix into Seurat object...')
  mex.trio[[1]] |>
    RowMergeSparseMatrices(mex.trio[-1]) |>
    CreateSeuratObject(min.cells = 3, min.features = 200)
}

DotPlot2d <- function(obj, feature, group.x, group.y, scaled = TRUE) {
  g1 <- obj |>
    mutate(subgroup = str_c({{ group.x }}, "bestsep", {{ group.y }})) |>
    DotPlot(feature, group.by = "subgroup") |>
    pluck("data") |>
    separate(id, sep = "bestsep", into = c("group.x", "group.y")) |>
    ggplot(aes(.data$group.x, .data$group.y, size = pct.exp)) +
    theme_pubr(legend = "right") +
    labs(color = "Average\nexpression", size = "Percent\nexpressed",
         title = feature)
  
  if (scaled) {
    g1 + geom_point(aes(color = avg.exp.scaled)) +
      scale_color_distiller(palette = "RdYlBu",
                            values = pretty_distiller(g1$data$avg.exp.scaled))
  } else {
    g1 + geom_point(aes(color = avg.exp)) +
      scale_color_distiller(palette = 'Reds', direction = 1)
  }
}

FindMarkersAcrossVar <- function(obj, split.by, group.by,
                                 ident.1, ident.2 = NULL,
                                 features = NULL,
                                 logfc.threshold = .1,
                                 min.pct = .01, fct_cluster = FALSE) {
  
  Idents(obj) <- split.by
  
  split.list <- Idents(obj) |> unique()
  
  stopifnot("split.by must have length > 1 in object" = length(split.list) > 1)
  
  TidyMarkers <- function(x) {
    FindMarkers(obj, subset.ident = x, group.by = group.by,
                  ident.1 = ident.1, ident.2 = ident.2,
                  features = features,
                  logfc.threshold = logfc.threshold,
                  min.pct = min.pct) |>
    as_tibble(rownames = 'gene') |>
    mutate(cluster = x)}
  
  cluster_deg <- split.list |>
    map(safely(TidyMarkers, quiet = F), .progress = T) |>
    list_transpose(simplify = FALSE) |>
    pluck('result') |>
    list_rbind()
  
  if (!fct_cluster) {
    cluster_deg <- cluster_deg |>
      mutate(cluster = as.character(cluster))
  }
  
  return(cluster_deg)
}

nose <- function(mat) {mat[1:5,1:5]}

#' Do slingshot pseudotime analysis on Seurat object
#'
#' @param sobj Seurat object to be analyzed.
#' @param reduction Reduction embedding for use. 'umap' for default.
#' @param cluster Cluster information for use. 'seurat_clusters' for default.
#' @param lineage Lineage names to show. Show all found lineages by default.
#'
#' @returns A modified `sobj` with lineages-pseudotime in `metadata` and curve coordinates in `misc`.
#' @export
#'
#' @examples
#' pbmc_small <- pbmc_small |> slingshot_seurat()
#' 
slingshot_seurat <- function(sobj, reduction = 'umap', 
                             cluster = 'seurat_clusters', lineage = NULL) {
  require(slingshot)
  embed_tb <- sobj |>
    Embeddings(reduction = reduction) 
  message('slingshot analysis...')
  sling_obj <- embed_tb |>
    slingshot(sobj@meta.data[[cluster]])
  
  curve_data <- sling_obj |> slingCurves() |> map('s')
  Misc(sobj, 'slingshot_curves') <- curve_data
  
  ptime <- sling_obj |>
    slingPseudotime()
  
  if (is_null(lineage)) {
    message('use all found lineages...')
    lineage_name <- colnames(ptime)
  } else {
    message(str_glue('use {lineage}...'))
    lineage_name <- lineage
    ptime <- ptime |>
      as.data.frame() |>
      select(all_of(lineage_name))
  }
  
  sobj <- sobj |>
    AddMetaData(ptime)
  
  g1 <- sobj |>
    FeaturePlot(lineage_name)
  
  if (length(lineage_name) == 1) {
    message('add curve for single lineage...')
    
    g1 <- g1 +
      geom_path(data = curve_data[[lineage_name]], aes(x = umap_1, y = umap_2),
                arrow = arrow(type = 'closed', length = unit(.1, 'inches')))
  }
  
  print(g1)
  
  return(sobj)
}

my36colors <-c('#E5D2DD', '#53A85F', '#F1BB72', '#F3B1A0', '#D6E7A3', '#57C3F3', '#476D87',
               '#E95C59', '#E59CC4', '#AB3282', '#23452F', '#BD956A', '#8C549C', '#585658',
               '#9FA3A8', '#E0D4CA', '#5F3D69', '#C5DEBA', '#58A4C3', '#E4C755', '#F7F398',
               '#AA9A59', '#E63863', '#E39A35', '#C1E6F3', '#6778AE', '#91D0BE', '#B53E2B',
               '#712820', '#DCC1DD', '#CCE0F5',  '#CCC9E6', '#625D9E', '#68A180', '#3A6963',
               '#968175')

pbmc_markers <- list(
  'B' = c('MS4A1','CD79A','CD19'),
  'PC' = c('JCHAIN'),
  'CD4' = c('IL7R','CCR7','S100A4'),
  'CD8' = 'CD8A',
  'NK' = c('GNLY','NKG7'),
  'Monocyte' = c('CD14','LYZ','FCGR3A','MS4A7'),
  'DC' = c('FCER1A','CST3'),
  'Mast' = c('TPSAB1','CPA3'),
  'pDC' = c('TLR7','TLR9')
)
